from functools import partial

import tensorflow as tf
from tensorflow.compat.v1 import logging
from tensorflow.compat.v1.keras.initializers import he_uniform
from tensorflow.keras.layers import (
    Concatenate,
    Conv2D,
    Conv2DTranspose,
    MaxPool2D,
    BatchNormalization,
    ELU,
    LeakyReLU,
    ReLU
)


def _get_conv_activation_layer(params):
    """

    :param params:
    :returns: Required Activation function.
    """
    conv_activation = params.get('conv_activation')
    if conv_activation == 'ReLU':
        return ReLU()
    elif conv_activation == 'ELU':
        return ELU()
    return LeakyReLU(0.2)


def _get_deconv_activation_layer(params):
    """

    :param params:
    :returns: Required Activation function.
    """
    deconv_activation = params.get('deconv_activation')
    if deconv_activation == 'LeakyReLU':
        return LeakyReLU(0.2)
    elif deconv_activation == 'ELU':
        return ELU()
    return ReLU()


def apply_image_unet(input_tensor,
                     output_name='image_unet',
                     params=None,
                     output_mask_logit=False):
    if params is None:
        params = {}
    logging.info(f'Apply unet for {output_name}')
    # conv_n_filters = params.get('conv_n_filters', [64, 128, 256, 512, 1024])
    conv_n_filters = params.get('conv_n_filters', [16, 32, 64, 128, 256, 512])
    num_classes = params.get("num_classes", 1)
    conv_activation_layer = _get_conv_activation_layer(params)
    deconv_activation_layer = _get_deconv_activation_layer(params)
    kernel_initializer = he_uniform(50)

    conv2d_factory = partial(
        Conv2D,
        kernel_size=(3, 3),
        strides=(1, 1),
        padding='same',
        kernel_initializer=kernel_initializer
    )

    max_pool_2d_factory = partial(
        MaxPool2D,
        pool_size=(2, 2)
    )

    # 收缩路径
    conv_down_1 = conv2d_factory(conv_n_filters[0])(input_tensor)
    conv_down_1 = BatchNormalization(axis=-1)(conv_down_1)
    conv_down_1 = conv_activation_layer(conv_down_1)
    conv_down_1 = conv2d_factory(conv_n_filters[0])(conv_down_1)
    conv_down_1 = BatchNormalization(axis=-1)(conv_down_1)
    conv_down_1 = conv_activation_layer(conv_down_1)
    pool_down_1 = max_pool_2d_factory()(conv_down_1)

    conv_down_2 = conv2d_factory(conv_n_filters[1])(pool_down_1)
    conv_down_2 = BatchNormalization(axis=-1)(conv_down_2)
    conv_down_2 = conv_activation_layer(conv_down_2)
    conv_down_2 = conv2d_factory(conv_n_filters[1])(conv_down_2)
    conv_down_2 = BatchNormalization(axis=-1)(conv_down_2)
    conv_down_2 = conv_activation_layer(conv_down_2)
    pool_down_2 = max_pool_2d_factory()(conv_down_2)

    conv_down_3 = conv2d_factory(conv_n_filters[2])(pool_down_2)
    conv_down_3 = BatchNormalization(axis=-1)(conv_down_3)
    conv_down_3 = conv_activation_layer(conv_down_3)
    conv_down_3 = conv2d_factory(conv_n_filters[2])(conv_down_3)
    conv_down_3 = BatchNormalization(axis=-1)(conv_down_3)
    conv_down_3 = conv_activation_layer(conv_down_3)
    pool_down_3 = max_pool_2d_factory()(conv_down_3)

    conv_down_4 = conv2d_factory(conv_n_filters[3])(pool_down_3)
    conv_down_4 = BatchNormalization(axis=-1)(conv_down_4)
    conv_down_4 = conv_activation_layer(conv_down_4)
    conv_down_4 = conv2d_factory(conv_n_filters[3])(conv_down_4)
    conv_down_4 = BatchNormalization(axis=-1)(conv_down_4)
    conv_down_4 = conv_activation_layer(conv_down_4)
    pool_down_4 = max_pool_2d_factory()(conv_down_4)

    conv_down_5 = conv2d_factory(conv_n_filters[4])(pool_down_4)
    conv_down_5 = BatchNormalization(axis=-1)(conv_down_5)
    conv_down_5 = conv_activation_layer(conv_down_5)
    conv_down_5 = conv2d_factory(conv_n_filters[4])(conv_down_5)
    conv_down_5 = BatchNormalization(axis=-1)(conv_down_5)
    conv_down_5 = conv_activation_layer(conv_down_5)

    deconv2d_factory = partial(
        Conv2DTranspose,
        kernel_size=(2, 2),
        strides=(2, 2),
        padding='same',
        kernel_initializer=kernel_initializer
    )

    crop = partial(tf.image.resize_with_crop_or_pad)
    # 扩展路径

    deconv_up_4 = deconv2d_factory(conv_n_filters[3])(conv_down_5)
    deconv_up_4 = BatchNormalization(axis=-1)(deconv_up_4)
    # crop_up_4 = crop(conv_down_4, tf.shape(deconv_up_4)[1], tf.shape(deconv_up_4)[2])
    # crop_up_4 = tf.reshape(crop_up_4, tf.shape(deconv_up_4))
    merge_up_4 = Concatenate(axis=-1)([deconv_up_4, conv_down_4])
    conv_up_4 = conv2d_factory(conv_n_filters[3])(merge_up_4)
    conv_up_4 = deconv_activation_layer(conv_up_4)
    conv_up_4 = conv2d_factory(conv_n_filters[3])(conv_up_4)
    conv_up_4 = deconv_activation_layer(conv_up_4)

    deconv_up_3 = deconv2d_factory(conv_n_filters[2])(conv_up_4)
    deconv_up_3 = BatchNormalization(axis=-1)(deconv_up_3)
    # crop_up_3 = crop(conv_down_3, tf.shape(deconv_up_3)[1], tf.shape(deconv_up_3)[2])
    # crop_up_3 = tf.reshape(crop_up_3, tf.shape(deconv_up_3))
    merge_up_3 = Concatenate(axis=-1)([deconv_up_3, conv_down_3])
    conv_up_3 = conv2d_factory(conv_n_filters[2])(merge_up_3)
    conv_up_3 = deconv_activation_layer(conv_up_3)
    conv_up_3 = conv2d_factory(conv_n_filters[2])(conv_up_3)
    conv_up_3 = deconv_activation_layer(conv_up_3)

    deconv_up_2 = deconv2d_factory(conv_n_filters[1])(conv_up_3)
    deconv_up_2 = BatchNormalization(axis=-1)(deconv_up_2)
    # crop_up_2 = crop(conv_down_2, tf.shape(deconv_up_2)[1], tf.shape(deconv_up_2)[2])
    # crop_up_2 = tf.reshape(crop_up_2, tf.shape(deconv_up_2))
    merge_up_2 = Concatenate(axis=-1)([deconv_up_2, conv_down_2])
    conv_up_2 = conv2d_factory(conv_n_filters[1])(merge_up_2)
    conv_up_2 = deconv_activation_layer(conv_up_2)
    conv_up_2 = conv2d_factory(conv_n_filters[1])(conv_up_2)
    conv_up_2 = deconv_activation_layer(conv_up_2)

    deconv_up_1 = deconv2d_factory(conv_n_filters[0])(conv_up_2)
    deconv_up_1 = BatchNormalization(axis=-1)(deconv_up_1)
    # crop_up_1 = crop(conv_down_1, tf.shape(deconv_up_1)[1], tf.shape(deconv_up_1)[2])
    # crop_up_1 = tf.reshape(crop_up_1, tf.shape(deconv_up_1))
    merge_up_1 = Concatenate(axis=-1)([deconv_up_1, conv_down_1])
    conv_up_1 = conv2d_factory(conv_n_filters[0])(merge_up_1)
    conv_up_1 = deconv_activation_layer(conv_up_1)
    conv_up_1 = conv2d_factory(conv_n_filters[0])(conv_up_1)
    conv_up_1 = deconv_activation_layer(conv_up_1)
    if not output_mask_logit:
        output = conv2d_factory(filters=num_classes, activation='sigmoid', kernel_size=(1, 1))(conv_up_1)
        return output
    return conv2d_factory(filters=num_classes, activation='softmax', kernel_size=(1, 1))(conv_up_1)


def image_unet(input_tensor, params=None):
    """ Model function applier. """
    return apply_image_unet(input_tensor, params=params)


def softmax_image_unet(input_tensor, params=None):
    return apply_image_unet(input_tensor, params=params, output_mask_logit=True)
